bert模型部署初探

发表于   |   分类于

Bert 模型部署

最近在研究google的bert相关知识，使用fine-tune方法训练了一个模型，但是看效果比较麻烦，每次都要重新load模型导致效率非常低下。
因此，参考网上以及同事的思路，设计了一套通用的bert model查看效果的工具。

主要内容

  1. 如何产出模型文件？

    1
    2
    estimator._export_to_tpu = False
    estimator.export_savedmodel("./export/", serving_input_fn)

  2. 如何避免重复加载？

    我们通过graph建立了sess，启动一个web应用，将sess保存在内存中，请求来的时候，直接使用sess进行predict，export_dir 是模型的路径。

    1
    2
    3
    4
    from tensorflow.python.saved_model import tag_constants
    self.graph_predict = tf.Graph()
    self.sess = tf.Session(graph = self.graph_predict)
    tf.saved_model.loader.load(self.sess, [tag_constants.SERVING], export_dir)

  3. 如何结合web框架？

    本次使用的是flask，为什么要使用flask，因为tf使用的是python，使用flask的话，可以直接将sess的实例保存在内存中，每次请求来的时候进行预测。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
from predictor import Predictor
one_predictor = Predictor(args)

app = Flask(__name__)
@app.route('/')
def hello_world():
    return 'Hello World!'

@app.route('/test', methods=['GET', 'POST'])
def index():
    if request.method == 'POST':
        query = (request.form.get('se_query').strip())
        result = one_predictor.predict_query(query)
        return render_template('test.html', result)
    
if __name__ == '__main__':
    server = pywsgi.WSGIServer(('0.0.0.0', 8001), app)
    server.serve_forever()

总结

使用flask封装bert fine-tune后的模型，极大的提高了策略调研的效率，还是要多思考，多尝试，多看tf的api相关文档。

相关代码

  1. https://github.com/c00h00g/bert_app_serving
------ 本文结束 ------
k